// Copyright 2014 Timothy Edward Baldwin

// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at

//     http://www.apache.org/licenses/LICENSE-2.0

// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.


#define __GNU_SOURCE
#include "navalis.h"

#include <assert.h>
#include <fcntl.h>
#include <linux/icmpv6.h>
#include <netdb.h>
#include <popt.h>
#include <signal.h>
#include <stdlib.h>
#include <sys/epoll.h>
#include <sys/signalfd.h>
#include <sys/socket.h>
#include <sys/stat.h>
#include <sys/timerfd.h>
#include <sys/types.h>
#include <unistd.h>

#include <fstream>
#include <iostream>
#include <map>
#include <sstream>
#include <vector>
using std::cout;
using std::cerr;
using std::endl;

typedef void (*event_function_t)();
typedef uint64_t hires_time;

struct trailer {
  uint32_t address;
  uint64_t token;
} __attribute__((packed));

struct remote_peer {
  uint64_t nounce_relay_discover = 0;
  uint64_t nounce_relay_local = 0;
  uint64_t nounce_relay_remote = 0;
  hires_time relay_discover_next = 0;
  hires_time relay_discover_interval = 1;
  hires_time relay_discover_deadline = 0;
  trailer remote_relay_temp;
  trailer remote_relay;
  uint8_t probe_time_recieved = 0;
};

struct local_peer {
  uint64_t nounce = 0;
  hires_time test_next = 0;
  hires_time test_interval = 1;
  hires_time test_deadline = 0;
  uint8_t initial_packets = 0;
};


int rand_fd;
int bubble_timer;
int bubble_deadline_timer;
int sock;
int rawsock;
int tun;
hires_time now;
int sig_fd;
const char *interface_name = "6a44";
std::map<ipv6_addr, local_peer> local_peers;
std::map<ipv6_addr, remote_peer> remote_peers;
ipv6_addr our_address;
uint64_t bubble_nounce;
uint32_t relay_address;
unsigned bubble_interval = 3;
bool have_address = false;
bool bubble_sent = false;
bool quit_pending = false;
bool script_running = false;
const char *script_path = PREFIX "/sbin/navalis-script";
int default_route_metric = 1028;
int current_route_metric;

uint32_t ip_checksum(uint32_t acc, const void *vdata, size_t length) {
  const char *data = reinterpret_cast<const char *>(vdata);
  for (size_t i = 0; i < length; i += 2) {
    uint16_t word;
    memcpy(&word, data + i, 2);
    acc += ntohs(word);
  }
  return acc;
}

void icmp_checksum(ipv6_packet &packet) {
  packet.icmp_checksum = 0;
  uint32_t acc = 0xFFFF;
  acc = ip_checksum(acc,&packet.source, 32);
  acc = ip_checksum(acc,&packet.length, 2);
  acc += packet.next_header;
  acc = ip_checksum(acc,&packet.icmp_type, sizeof(packet) - 40);
  packet.icmp_checksum = htons(~(acc + (acc >> 16)));
}

static void set_timer(int fd, long milliseconds) {
  struct itimerspec spec;
  set_zero(spec);
  spec.it_value.tv_sec = milliseconds / 1000;
  spec.it_value.tv_nsec = (milliseconds % 1000) * 1000000;
  timerfd_settime(fd, 0, &spec, nullptr);
}

hires_time add_jitter(unsigned t) {
  uint32_t r;
  set_random(r);
  return t * 1000 - ((static_cast<uint64_t>(r) * t) >> 23);
}

void display_addr(const ipv6_addr& n) {
  char ip6str[INET6_ADDRSTRLEN];
  inet_ntop(AF_INET6, &n, ip6str, INET6_ADDRSTRLEN);
  cout << ip6str << endl;
}

void run_process(const char * const *args) {
  int pid = fork();
  assert(pid >= 0);

  if (!pid) {
    execv(args[0], const_cast<char *const *>(args));
    _exit(1);
  }
  script_running = true;
}

void run_script() {

  static bool script_have_address = false;
  static ipv6_addr script_address;
  static int script_route_metric;

  // Don't run multiple script instances.
  if (script_running) return;

  char ip6str[INET6_ADDRSTRLEN];
  char metric[11];
  const char *args[6];
  args[0] = script_path;
  args[2] = interface_name;
  args[3] = ip6str;
  args[4] = metric;
  args[5] = nullptr;

  if (script_have_address && (!have_address || script_address != our_address)) {
    inet_ntop(AF_INET6, &script_address, ip6str, INET6_ADDRSTRLEN);
    args[1] = "remove";
    sprintf(metric, "%i", script_route_metric);
    run_process(args);
    script_have_address = false;

  } else if (have_address && !script_have_address) {
    inet_ntop(AF_INET6, &our_address, ip6str, INET6_ADDRSTRLEN);
    args[1] = (script_address == our_address) ? "readd" : "add";
    sprintf(metric, "%i", current_route_metric);
    run_process(args);
    script_address = our_address;
    script_have_address = true;
    script_route_metric = current_route_metric;

  } else if (script_have_address && script_route_metric != current_route_metric) {
    args[1] = "change-metric";
    sprintf(ip6str, "%i", script_route_metric);
    sprintf(metric, "%i", current_route_metric);
    run_process(args);
    script_route_metric = current_route_metric;

  } else if (quit_pending) {
    exit(0);
  }
}

void construct_echo_request(ipv6_packet& packet) {
  set_zero(packet);
  packet.version = 6;
  packet.hoplimit = 100;
  packet.next_header = IPPROTO_ICMPV6;
  packet.length = htons(sizeof(packet) - 40);
  packet.source = our_address;
  packet.icmp_type = ICMPV6_ECHO_REQUEST;
}

void send_bubble() {

  // Simple rate limit
  if (bubble_sent) return;

  // Send bubble to 6a44 relay
  bubble_t bubble;
  set_zero(bubble.prefix);
  struct sockaddr_in saddr;
  saddr.sin_family = AF_INET;
  saddr.sin_port = PORT_6a44;
  saddr.sin_addr.s_addr = relay_address;
  bubble.nounce = bubble_nounce;
  sendto(sock, &bubble, sizeof(bubble), 0, sa_cast(&saddr), sizeof(saddr));
  bubble_sent = true;
}

void remove_address() {
  // Remove address from interface if present
  have_address = false;
  run_script();
}

void process_bubble(bubble_t bubble) {

  if (quit_pending) return;

  // Reset timers
  bubble_interval = 2;
  set_timer(bubble_timer, add_jitter(30));
  set_timer(bubble_deadline_timer, add_jitter(120));

  // Was an address provided?
  if (bubble.address.prefix) {

    // If change of address flush nieghbour cache.
    if (our_address != bubble.address) {
      local_peers.clear();
      remote_peers.clear();
    }

    if (our_address.private_ipv4 != bubble.address.private_ipv4) {
      // Bind proto41 socket
      struct sockaddr_in saddr;
      saddr.sin_family = AF_INET;
      saddr.sin_port = 0;
      saddr.sin_addr.s_addr = bubble.address.private_ipv4;
      bind(rawsock, sa_cast(&saddr), sizeof(saddr));
    }

    // Save address
    have_address = true;
    our_address = bubble.address;
    current_route_metric = default_route_metric;

  } else {
    // No address provided
    have_address = false;
  }

  // Send address to script
  run_script();

  // Change nounce
  set_random(bubble_nounce);
}

void retransmit_bubble() {

  // Check if maxinum retransmission interval reached.
  if (bubble_interval > 300) {
    bubble_interval = 300;
    set_random(bubble_nounce); // Prevent attacks with old packet captures.
  }

  // Send bubble
  send_bubble();
  bubble_sent = false;

  // Reset timer
  set_timer(bubble_timer, add_jitter(bubble_interval));
  bubble_interval *= 2;
}

void recv_ipv6(ipv6_packet &packet, size_t size) {

  // Forward to IPv6 stack.
  write(tun, &packet, packet.size());

  // Is it a ICMP Echo Reply?
  if (packet.next_header != IPPROTO_ICMPV6
      || packet.icmp_type != ICMPV6_ECHO_REPLY
      || packet.nounce == 0)
    return;

  // Is it a response from proto41 test?
  if (packet.source.prefix == our_address.prefix) {
    auto peer = local_peers.find(packet.source);
    if (peer != local_peers.end() && packet.nounce == peer->second.nounce) {
      peer->second.test_next = now + add_jitter(900);
      peer->second.test_deadline = now + add_jitter(1800);
      peer->second.test_interval = 1;
      peer->second.initial_packets = 0;
      peer->second.nounce = 0;
    }
  }

  if (packet.echo_magic != ECHO_MAGIC) return;
  auto peer_iterator = remote_peers.find(packet.source);
  if (peer_iterator == remote_peers.end()) return;
  remote_peer *peer = &peer_iterator->second;

  if (packet.nounce == peer->nounce_relay_discover) {
    if (size >= packet.size() + sizeof(trailer)) {
      peer->nounce_relay_discover = 0;

      // Read trailer.
      memcpy(&peer->remote_relay_temp, u8_cast(&packet) + packet.size(), sizeof(trailer));

      // Construct probe packet
      ipv6_packet probe;
      construct_echo_request(probe);
      probe.echo_magic = ECHO_MAGIC;
      probe.source.private_ipv4 = 0;
      probe.destination = packet.source;
      struct sockaddr_in saddr;
      saddr.sin_family = AF_INET;
      saddr.sin_port = PORT_6a44;

      // Send probe via local relay
      set_random(peer->nounce_relay_local);
      probe.nounce = peer->nounce_relay_local;
      icmp_checksum(probe);
      saddr.sin_addr.s_addr = relay_address;
      sendto(sock, &probe, sizeof(probe), 0, sa_cast(&saddr), sizeof(saddr));

      // Send probe via remote relay
      set_random(peer->nounce_relay_remote);
      probe.nounce = peer->nounce_relay_remote;
      icmp_checksum(probe);
      saddr.sin_addr.s_addr = peer->remote_relay_temp.address;

      struct iovec iov[2];
      iov[0].iov_base = &probe;
      iov[0].iov_len = sizeof(probe);
      iov[1].iov_base = &peer->remote_relay_temp.token;
      iov[1].iov_len = sizeof(trailer::token);

      struct msghdr mh;
      mh.msg_name = &saddr;
      mh.msg_namelen = sizeof(saddr);
      mh.msg_iov = iov;
      mh.msg_iovlen = 2;
      mh.msg_control = nullptr;
      mh.msg_controllen = 0;
      mh.msg_flags = 0;

      sendmsg(sock, &mh, 0);

    } else {
      // Remote relay does not provide it's unicast IPv4 address so retry in 1 hour.
      peer->relay_discover_deadline = 0;
      peer->relay_discover_next = now + add_jitter(3600);
      peer->relay_discover_interval = 60;
      peer->nounce_relay_local = 0;
      peer->nounce_relay_remote = 0;
    }

  } else if (packet.nounce == peer->nounce_relay_local) {
    if (peer->probe_time_recieved == 2) {
      // Response via remote relay before local
      peer->relay_discover_deadline = now + 360 * 1000;
      peer->relay_discover_next = now + add_jitter(300);
      peer->relay_discover_interval = 5;
      peer->remote_relay = peer->remote_relay_temp;
    }
    peer->probe_time_recieved |= 1;

  } else if (packet.nounce == peer->nounce_relay_remote) {
    if (peer->probe_time_recieved == 1) {
      // Response via local relay before remote
      peer->relay_discover_deadline = 0;
      peer->relay_discover_next = now + add_jitter(600);
      peer->relay_discover_interval = 10;
    }
    peer->probe_time_recieved |= 2;
  }
}

inline bool proto41_ok(const ipv6_packet &packet) {
  return packet.destination.prefix == packet.source.prefix
      && (packet.destination.public_ipv4 == packet.source.public_ipv4
          || (packet.destination.public_ipv4 == packet.destination.private_ipv4
              && packet.source.public_ipv4 == packet.source.private_ipv4));
}

void send_ipv6() {
  // Read packet from tun device
  udp_packet packet;
  ssize_t size = read(tun, &packet, sizeof(packet));

  // Validate packet
  if (packet.v6.version != 6 || size != packet.v6.size()) return;
  if (!same_client(packet.v6.source, our_address)) return; // FIXME return ICMPv6 error

  // Set up address struct
  struct sockaddr_in saddr;
  saddr.sin_family = AF_INET;
  saddr.sin_port = PORT_6a44;

  // A special case...
  if (packet.v6.next_header == IPPROTO_ICMPV6
      && packet.v6.icmp_type == ICMPV6_ECHO_REPLY
      && packet.v6.echo_magic == ECHO_MAGIC) {

    // This a probe response, so always send via local relay for consitent timing.
    saddr.sin_addr.s_addr = relay_address;
    sendto(sock, &packet, size, 0, sa_cast(&saddr), sizeof(saddr));
    return;
  }

  // Can we send it direct?
  if (proto41_ok(packet.v6) && packet.v6.source.private_ipv4 == our_address.private_ipv4) {

    local_peer *peer = &local_peers[packet.v6.destination];
    saddr.sin_addr.s_addr = packet.v6.destination.private_ipv4;

    // Send another test packet?
    if (now > peer->test_next) {
      ipv6_packet probe;
      set_random(peer->nounce);
      construct_echo_request(probe);
      probe.nounce = peer->nounce;
      probe.destination = packet.v6.destination;
      icmp_checksum(probe);
      sendto(rawsock, &probe, sizeof(probe), 0, sa_cast(&saddr), sizeof(saddr));

      // Exponential backoff retransmissions upto hourly.
      peer->test_next = now + add_jitter(peer->test_interval);
      peer->test_interval *=2;
      if (peer->test_interval > 3600) peer->test_interval = 3600;
    }

    // Good direct path? (No NATs)
    if (now < peer->test_deadline) {
      sendto(rawsock, &packet, size, 0, sa_cast(&saddr), sizeof(saddr));
      return;
    }

    // Send first 4 packets both ways for fast and reliable setup.
    if (peer->initial_packets < 4) {
      ++peer->initial_packets;
      sendto(rawsock, &packet, size, 0, sa_cast(&saddr), sizeof(saddr));
      saddr.sin_addr.s_addr = relay_address;
      sendto(sock, &packet, size, 0, sa_cast(&saddr), sizeof(saddr));
      return;
    }
  }

  remote_peer *peer = &remote_peers[packet.v6.destination];
  if (peer->relay_discover_next < now) {

    // Send ICMP Echo in order to discover relay address
    ipv6_packet probe;
    set_random(peer->nounce_relay_discover);
    construct_echo_request(probe);
    probe.echo_magic = ECHO_MAGIC;
    probe.source.private_ipv4 = 0;
    probe.nounce = peer->nounce_relay_discover;
    probe.destination = packet.v6.destination;
    icmp_checksum(probe);
    saddr.sin_addr.s_addr = relay_address;
    sendto(sock, &probe, sizeof(probe), 0, sa_cast(&saddr), sizeof(saddr));

    // Reset state of received responses
    peer->remote_relay_temp.address = 0;
    peer->probe_time_recieved = 0;

    // Exponential backoff retransmissions upto hourly.
    peer->relay_discover_next = now + add_jitter(peer->relay_discover_interval);
    peer->relay_discover_interval *= 2;
    if (peer->relay_discover_interval > 3600) peer->relay_discover_interval = 3600;
  }

  if (now < peer->relay_discover_deadline) {
    saddr.sin_addr.s_addr = peer->remote_relay.address;

    struct iovec iov[2];
    iov[0].iov_base = &packet;
    iov[0].iov_len = size;
    iov[1].iov_base = &peer->remote_relay.token;
    iov[1].iov_len = sizeof(trailer::token);

    struct msghdr mh;
    mh.msg_name = &saddr;
    mh.msg_namelen = sizeof(saddr);
    mh.msg_iov = iov;
    mh.msg_iovlen = 2;
    mh.msg_control = nullptr;
    mh.msg_controllen = 0;
    mh.msg_flags = 0;

    sendmsg(sock, &mh, 0);
  } else {
    saddr.sin_addr.s_addr = relay_address;
    sendto(sock, &packet, size, 0, sa_cast(&saddr), sizeof(saddr));
  }
}

void recv_ipv4_udp() {
  struct sockaddr_in saddr;
  udp_packet packet;
  uint64_t cmsg_space[8];

  struct iovec iov;
  iov.iov_base = &packet;
  iov.iov_len = sizeof(packet);

  struct msghdr mh;
  mh.msg_name = &saddr;
  mh.msg_namelen = sizeof(saddr);
  mh.msg_iov = &iov;
  mh.msg_iovlen = 1;
  mh.msg_control = cmsg_space;
  mh.msg_controllen = sizeof(cmsg_space);
  mh.msg_flags = 0;

  ssize_t size = recvmsg(sock, &mh, 0);

  // Discard packets not from relay.
  if (saddr.sin_port != PORT_6a44 || relay_address != saddr.sin_addr.s_addr || size < 20) return;


  if (size < 40 && packet.bubble.nounce == bubble_nounce) {
    // Bubble - change address if needed

    // Extract local IPv4 address
    struct cmsghdr *cmsg = CMSG_FIRSTHDR(&mh);
    while(cmsg->cmsg_level != IPPROTO_IP || cmsg->cmsg_type != IP_PKTINFO)
      cmsg = CMSG_NXTHDR(&mh, cmsg);
    struct in_pktinfo *pi = reinterpret_cast<struct in_pktinfo *>(CMSG_DATA(cmsg));
    packet.bubble.address.private_ipv4 = pi->ipi_spec_dst.s_addr;

    process_bubble(packet.bubble);

  } else if (packet.v6.version != 6 || size < packet.v6.size()) {
    // Do nothing - not a valid packet

  } else if (same_client(packet.v6.destination, our_address)) {
    // Packet is for us.
    // Then process
    recv_ipv6(packet.v6, size);

  } else {
    // Send bubble - address suspect.
    send_bubble();
  }

}

void recv_ipv4_proto41() {
  union {
    uint8_t bytes[1500];
    struct iphdr header;
  } ipv4_packet __attribute__ ((aligned(4)));

  ssize_t size = read(rawsock, &ipv4_packet, sizeof(ipv4_packet));
  if (size < static_cast<ssize_t>(sizeof(struct iphdr))) return;

  ipv6_packet *packet = reinterpret_cast<struct ipv6_packet*>(ipv4_packet.bytes + ipv4_packet.header.ihl * 4);
  size -= ipv4_packet.header.ihl * 4;
  if (packet->version != 6 || size < packet->size()) return;

  if (proto41_ok(*packet)
      && same_client(packet->destination, our_address)
      && packet->source.private_ipv4 == ipv4_packet.header.saddr
      && packet->destination.private_ipv4 == ipv4_packet.header.daddr)
    recv_ipv6(*packet, size);
}

inline int epoll_add(int epfd, int fd, event_function_t f, uint32_t events = EPOLLIN) {
  struct epoll_event event;
  event.data.ptr = reinterpret_cast<void *>(f);
  event.events = events;
  return epoll_ctl(epfd, EPOLL_CTL_ADD, fd, &event);
}

void recv_signalfd() {

  // Retrive pending signal
  struct signalfd_siginfo sig;
  ssize_t s = read(sig_fd, &sig, sizeof(sig));
  if (s != sizeof(sig)) return;

  // And process
  switch(sig.ssi_signo) {
  case SIGCHLD:
    script_running = false;
    run_script();
    break;
  case SIGUSR1:
    local_peers.clear();
    remote_peers.clear();
  case SIGHUP:
    bubble_interval = 2;
    retransmit_bubble();
    set_timer(bubble_deadline_timer, add_jitter(4));
    break;
  case SIGTERM:
    have_address = false;
    quit_pending = true;
    remove_address();
    break;
  }
}

namespace args {

const char *pidfile;
const char *relay;

struct poptOption options[] = {
{"pidfile", 0, POPT_ARG_STRING, &pidfile, 0, "PID File", nullptr},
{"relay", 'r', POPT_ARG_STRING, &relay, 0, "Relay IPv4 address or hostname (required)", nullptr},
{"interface-name", 'i', POPT_ARG_STRING, &interface_name, 0, "Interface name", nullptr},
{"script", 0, POPT_ARG_STRING, &script_path, 0, "Script", nullptr},
{"metric", 0, POPT_ARG_INT, &default_route_metric, 0, "Default route metric", nullptr},
POPT_AUTOHELP
POPT_TABLEEND
};
}

int main(int argc, const char **argv) {

  // Handle signals
  sigset_t sigs;
  sigemptyset(&sigs);
  sigaddset(&sigs, SIGCHLD);
  sigaddset(&sigs, SIGUSR1);
  sigaddset(&sigs, SIGHUP);
  sigaddset(&sigs, SIGTERM);
  sigprocmask(SIG_BLOCK, &sigs, nullptr);
  sig_fd = signalfd(-1, &sigs, SFD_NONBLOCK | SFD_CLOEXEC);

  // Read command line arguments
  poptContext opts = poptGetContext(nullptr, argc, argv, args::options, 0);
  poptGetNextOpt(opts);
  poptFreeContext(opts);

  // Open random number device file
  rand_fd = open("/dev/urandom", O_RDONLY | O_CLOEXEC);

  // Open timers
  bubble_timer = timerfd_create(CLOCK_MONOTONIC, TFD_NONBLOCK | TFD_CLOEXEC);
  bubble_deadline_timer = timerfd_create(CLOCK_MONOTONIC, TFD_NONBLOCK | TFD_CLOEXEC);

  // Open UDP socket
  sock = socket(AF_INET, SOCK_DGRAM | SOCK_NONBLOCK | SOCK_CLOEXEC , IPPROTO_UDP);
  int one = 1;
  setsockopt(sock, IPPROTO_IP, IP_PKTINFO, &one, sizeof(one));

  // Open Proto 41 socket
  rawsock = socket(AF_INET, SOCK_RAW | SOCK_NONBLOCK | SOCK_CLOEXEC , IPPROTO_IPV6);

  // Initialise bubble nounce
  set_random(bubble_nounce);

  // Read IP address from command line.
  struct addrinfo hints, *addresses;
  set_zero(hints);
  hints.ai_family = AF_INET;
  hints.ai_socktype = SOCK_DGRAM;
  hints.ai_protocol = IPPROTO_UDP;
  getaddrinfo(args::relay, "3544", &hints, &addresses);
  relay_address = sin_cast(addresses->ai_addr)->sin_addr.s_addr;
  freeaddrinfo(addresses);

  // Create TUN device
  struct ifreq ifr;
  tun = open("/dev/net/tun", O_RDWR | O_NONBLOCK | O_CLOEXEC);
  set_zero(ifr);
  ifr.ifr_flags = IFF_TUN | IFF_NO_PI;
  strncpy(ifr.ifr_name, interface_name, IFNAMSIZ);
  ioctl(tun, TUNSETIFF, &ifr);

  // Set TUN device MTU
  ifr.ifr_mtu = 1280;
  ioctl(sock, SIOCSIFMTU, &ifr);

  // Remove multicast
  ioctl(sock, SIOCGIFFLAGS, &ifr);
  ifr.ifr_flags &= ~IFF_MULTICAST;
  ioctl(sock, SIOCSIFFLAGS, &ifr);

  // Set up event loop
  int ep = epoll_create(4);
  epoll_add(ep, bubble_timer, retransmit_bubble, EPOLLIN | EPOLLET);
  epoll_add(ep, bubble_deadline_timer, remove_address, EPOLLIN | EPOLLET);
  epoll_add(ep, sock, recv_ipv4_udp);
  epoll_add(ep, rawsock, recv_ipv4_proto41);
  epoll_add(ep, tun, send_ipv6);
  epoll_add(ep, sig_fd, recv_signalfd);

  // Write pid file
  if (args::pidfile) std::ofstream(args::pidfile) << getpid() << endl;

  // Send inital bubble
  retransmit_bubble();

  while (true) {

    // Wait for event
    struct epoll_event event;
    epoll_wait(ep, &event, 1, -1);

    // Read current time
    struct timespec t;
    clock_gettime(CLOCK_MONOTONIC, &t);
    now = static_cast<hires_time>(t.tv_sec) * 1000 + t.tv_nsec /  1000000;

    // Call event routine
    reinterpret_cast<event_function_t>(event.data.ptr)();
  }
}
